/*
 * Copyright (c) 2016-present, Facebook, Inc.
 * All rights reserved.
 *
 * This source code is licensed under both the BSD-style license (found in the
 * LICENSE file in the root directory of this source tree) and the GPLv2 (found
 * in the COPYING file in the root directory of this source tree).
 */

/* Compression level or 0 to disable */
#define DO_ZLIB 9
/* Compression level or 0 to disable */
#define DO_ZSTD 0
/* Buffer size */
#define BUFFER_SIZE 4096

#include <linux/fs.h>
#include <linux/kernel.h>
#include <linux/module.h>
#include <linux/uaccess.h>
#include <linux/vmalloc.h>

#if DO_ZSTD
#include <linux/zstd.h>
#endif

#if DO_ZLIB
#include <linux/zlib.h>
#endif

/* Device name to pass to register_chrdev(). */
#define DEVICE_NAME "zstd_compress_test"

/* Dynamically allocated device major number */
static int device_major;

/*
 * We reuse the same state, and thus can compress only one file at a time.
 */
static bool device_is_open;

static void* workspace = NULL;

/*
 * Input buffer used to put data coming from userspace.
 */
static uint8_t buffer_in[BUFFER_SIZE];
static uint8_t buffer_out[BUFFER_SIZE];

static uint64_t uncompressed_len;
static uint64_t compressed_len;

#if DO_ZSTD

static ZSTD_CStream* state;

static ZSTD_inBuffer input = {
    .src = buffer_in,
    .size = sizeof(buffer_in),
    .pos = sizeof(buffer_in),
};

static ZSTD_outBuffer output = {
    .dst = buffer_out,
    .size = sizeof(buffer_out),
    .pos = sizeof(buffer_out),
};

#endif /* DO_ZSTD */

#if DO_ZLIB

static z_stream state = {
    .next_in = buffer_in,
    .avail_in = 0,
    .total_in = 0,

    .next_out = buffer_out,
    .avail_out = sizeof(buffer_out),
    .total_out = 0,

    .msg = NULL,
    .state = NULL,
    .workspace = NULL,
};

#endif /* DO_ZLIB */

static int zstd_compress_test_open(struct inode* i, struct file* f) {
    if (device_is_open)
        return -EBUSY;

    device_is_open = true;

    uncompressed_len = compressed_len = 0;

#if DO_ZSTD
    if (ZSTD_isError(ZSTD_resetCStream(state, 0)))
        return -EIO;
#endif

#if DO_ZLIB
    if (zlib_deflateReset(&state) != Z_OK)
        return -EIO;
#endif

    printk(KERN_INFO DEVICE_NAME ": opened\n");
    return 0;
}

static int zstd_compress_test_release(struct inode* i, struct file* f) {
    device_is_open = false;

#if DO_ZSTD
    do {
        size_t ret;

        output.pos = 0;
        ret = ZSTD_endStream(state, &output);
        if (ZSTD_isError(ret)) {
            printk(KERN_INFO DEVICE_NAME ": zstd end error %u\n", ZSTD_getErrorCode(ret));
            return -EIO;
        }
        compressed_len += output.pos;
    } while (output.pos != output.size);
#endif

#if DO_ZLIB
    for (;;) {
        int ret;

        state.next_out = buffer_out;
        state.avail_out = sizeof(buffer_out);
        ret = zlib_deflate(&state, Z_FINISH);
        compressed_len += sizeof(buffer_out) - state.avail_out;
        if (ret == Z_STREAM_END)
            break;
        if (ret != Z_OK) {
            printk(KERN_INFO DEVICE_NAME ": zlib end error %d: %s\n", ret, state.msg);
            return -EIO;
        }
    }
#endif

    printk(KERN_INFO DEVICE_NAME ": uncompressed_len = %llu\n", uncompressed_len);
    printk(KERN_INFO DEVICE_NAME ": compressed_len   = %llu\n", compressed_len);
    printk(KERN_INFO DEVICE_NAME ": closed\n");
    return 0;
}

/*
 * Hash the data given to us from userspace.
 */
static ssize_t zstd_compress_test_write(struct file* file, const char __user* buf, size_t size, loff_t* pos) {
    size_t remaining = size;

    while (remaining > 0) {
        size_t const copy_size = min(remaining, sizeof(buffer_in));

        if (copy_from_user(buffer_in, buf, copy_size))
            return -EFAULT;
        buf += copy_size;
        remaining -= copy_size;
        uncompressed_len += copy_size;

#if DO_ZSTD
        input.pos = 0;
        input.size = copy_size;
        while (input.pos != input.size) {
            size_t ret;

            output.pos = 0;
            ret = ZSTD_compressStream(state, &output, &input);
            if (ZSTD_isError(ret)) {
                printk(KERN_INFO DEVICE_NAME ": zstd compress error %u\n", ZSTD_getErrorCode(ret));
                return -EIO;
            }
            compressed_len += output.pos;
        }
#endif
#if DO_ZLIB
        state.next_in = buffer_in;
        state.avail_in = copy_size;
        while (state.avail_in > 0) {
            int ret;

            state.next_out = buffer_out;
            state.avail_out = sizeof(buffer_out);
            ret = zlib_deflate(&state, Z_NO_FLUSH);
            compressed_len += sizeof(buffer_out) - state.avail_out;
            if (ret != Z_OK) {
                printk(KERN_INFO DEVICE_NAME ": zlib end error %d: %s\n", ret, state.msg);
                return -EIO;
            }
        }
#endif
    }
    return size;
}
/* register the character device. */
static int __init zstd_compress_test_init(void) {
    static const struct file_operations fileops = {.owner = THIS_MODULE,
                                                   .open = &zstd_compress_test_open,
                                                   .release = &zstd_compress_test_release,
                                                   .write = &zstd_compress_test_write};
    size_t workspace_size = 0;
#if DO_ZSTD
    ZSTD_parameters params;
#endif

    device_major = register_chrdev(0, DEVICE_NAME, &fileops);
    if (device_major < 0) {
        return device_major;
    }

#if DO_ZSTD
    params = ZSTD_getParams(DO_ZSTD, 0, 0);
    workspace_size = ZSTD_CStreamWorkspaceBound(params.cParams);

    if (!(workspace = vmalloc(workspace_size)))
        goto fail;
    if (!(state = ZSTD_initCStream(params, 0, workspace, workspace_size)))
        goto fail;
#endif

#if DO_ZLIB
    workspace_size = zlib_deflate_workspacesize(MAX_WBITS, MAX_MEM_LEVEL);

    if (!(workspace = vmalloc(workspace_size)))
        goto fail;
    state.workspace = workspace;
    if (zlib_deflateInit(&state, DO_ZLIB) != Z_OK)
        goto fail;
#endif

    printk(KERN_INFO DEVICE_NAME ": module loaded\n");
    printk(KERN_INFO DEVICE_NAME ": compression requires %zu bytes of memory\n", workspace_size);
    printk(KERN_INFO DEVICE_NAME
           ": Create a device node with "
           "'mknod " DEVICE_NAME
           " c %d 0' and write data "
           "to it.\n",
           device_major);
    return 0;

fail:
    printk(KERN_INFO DEVICE_NAME ": failed to load module\n");
    if (workspace) {
        vfree(workspace);
        workspace = NULL;
    }
    return -ENOMEM;
}

static void __exit zstd_compress_test_exit(void) {
    unregister_chrdev(device_major, DEVICE_NAME);
#if DO_ZLIB
    zlib_deflateEnd(&state);
#endif
    if (workspace) {
        vfree(workspace);
        workspace = NULL;
    }
    printk(KERN_INFO DEVICE_NAME ": module unloaded\n");
}

module_init(zstd_compress_test_init);
module_exit(zstd_compress_test_exit);

MODULE_DESCRIPTION("Zstd compression tester");
MODULE_VERSION("1.0");

MODULE_LICENSE("Dual BSD/GPL");
